I've cleared the notebook output, but here's a rendered file from the most recent commit leifdenby@daa5f02:
saving_graphs_for_neural_lam.html

Creating these plots @joeloskarsson (see end of notebook) I think what is currently implemented when saving to neural-lam format actually has a bug, the one which you're fixing in https://github.com/mllam/weather-model-graphs/pull/34/files#diff-bdf567174bcb3d8da552f28644f860a91c5019a07b6238b7d85e35bd346c9edbL118. What do you think?

Adding link to comment @joeloskarsson made which I think is important here:

assumptions wrt node index order and subgraph order (in terms of m2g, g2m and m2m, and its levels) that neural-lam assumes?

• There is no assumption about the subgraph order, in the sense of m2g, g2m and m2m, as they are stored separately.

• In hierarchical mesh graphs the bottom level should come first, followed by higher levels in the hierarchy. But note that this is only the ordering of levels in the list, there are no requirements on one level having node indices higher/lower than any other.

• We are currently saying in neural-lam that mesh nodes must have the index before grid nodes (https://github.com/mllam/neural-lam/blob/4d8aceade489ff23f1a11cf0c9f9cd967902a35a/neural_lam/models/base_graph_model.py#L22), but this is not crucial and we should really change this as we integrate wmg better with neural-lam. I already did this on the research branch. See the point "Move all edge_index reindexing ..." in Enabling training on interior and boundary datastore neural-lam#138.

This should be very useful indeed! Looking at the output there seems to a bug yes. It could very well be the label sorting that's missing, that you linked to as fixed in #34. Have you tried rebasing this on top of that PR and see if you have the same issue?

Have you tried rebasing this on top of that PR and see if you have the same issue?

No, I will try that and see :) thanks

Now that #34 is in main (with the sorting that you implemented @joeloskarsson) I merged main into my branch for this work.

The edges in the loaded files (neural-lam pickled pytorch tensor files) was previously:

Now it looks like this:

And here the edges from the networkx.DiGraph as defined by the node labels:

So I think there is still something going wrong since the mesh indexes for the m2m subgraph edges start at zero (of course this might not matter currently in neural-lam since we always offset the indexes so that the first index is 0, and then assume that the index values are consecutive - that's right, right?)

I think a bit more digging is needed here...

I've included the cell outputs in the notebook in the most recent commit (a7a8e9e).

Ok, I think that maybe we viewed different things as wrong. Or maybe I just did not pay close enough attention. I was mostly worried that not all mesh nodes seemed to be part of an edge in g2m. But maybe that was just how you set up the graph?

If I understand you correctly, you view the problem as m2m being saved with the first mesh node having index 0 in edge_index? But I don't think this is something you can get away from, if converting from networkx to pyg using pyg.utils. There are no grid nodes in the m2m graph, so naturally the first mesh node will have index 0. I suppose you could offset all mesh node indices with num_grid_nodes in the tensor, but that seems like a bit of a hack. But why is this a problem? As you wrote, neural-lam anyhow needs the m2m node indices to start with 0, and if they don't we now reindex them there.

There are no grid nodes in the m2m graph, so naturally the first mesh node will have index 0. I suppose you could offset all mesh node indices with num_grid_nodes in the tensor, but that seems like a bit of a hack. But why is this a problem? As you wrote, neural-lam anyhow needs the m2m node indices to start with 0, and if they don't we now reindex them there.

I had expected that the node-indices in the tensors saved to disk would be globally unique integers. Rereading what you've written here makes me realise we should document as part of the spec in mllam/neural-lam#323 the node-index values in the adjacency-list tensors and how they relate to the indices of the feature tensors (in effect what we are discussing here: mllam/neural-lam#323 (comment))

I think it would make most sense if the way we reference nodes in the edge_index values was if each nodeset (e.g. grid or mesh nodes) would start with the first node being at index value 0 and the last at N_grid-1 (for grid nodes) or N_mesh-1 (for mesh nodes). That isn't what is happening in the image I posted above (middle one) where I have loaded the tensor files produced by wmg. In that image for example the g2m edges the node-index for the mesh (g2m dst) nodes starts at 12 rather than 0.